#!/usr/bin/env python

# script to call a DRS GetNCChanges from the command line
# this is useful for plugfest testing

import sys
from optparse import OptionParser

sys.path.insert(0, "bin/python")

import samba, ldb
import samba.getopt as options
from samba.dcerpc import drsuapi, misc
from samba.samdb import SamDB
from samba.auth import system_session
from samba.ndr import ndr_unpack

def do_DsBind(drs):
    '''make a DsBind call, returning the binding handle'''
    bind_info = drsuapi.DsBindInfoCtr()
    bind_info.length = 28
    bind_info.info = drsuapi.DsBindInfo28()
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_BASE
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_ASYNC_REPLICATION
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_REMOVEAPI
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_MOVEREQ_V2
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_GETCHG_COMPRESS
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_DCINFO_V1
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_RESTORE_USN_OPTIMIZATION
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_KCC_EXECUTE
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_ADDENTRY_V2
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_LINKED_VALUE_REPLICATION
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_DCINFO_V2
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_INSTANCE_TYPE_NOT_REQ_ON_MOD
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_CRYPTO_BIND
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_GET_REPL_INFO
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_STRONG_ENCRYPTION
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_DCINFO_V01
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_TRANSITIVE_MEMBERSHIP
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_ADD_SID_HISTORY
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_POST_BETA3
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_GET_MEMBERSHIPS2
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_GETCHGREQ_V6
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_NONDOMAIN_NCS
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_GETCHGREQ_V8
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_GETCHGREPLY_V5
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_GETCHGREPLY_V6
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_ADDENTRYREPLY_V3
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_GETCHGREPLY_V7
    bind_info.info.supported_extensions	|= drsuapi.DRSUAPI_SUPPORTED_EXTENSION_VERIFY_OBJECT
    (info, handle) = drs.DsBind(misc.GUID(drsuapi.DRSUAPI_DS_BIND_GUID), bind_info)
    return handle


def drs_get_rodc_partial_attribute_set(samdb):
    '''get a list of attributes for RODC replication'''
    partial_attribute_set = drsuapi.DsPartialAttributeSet()
    partial_attribute_set.version = 1

    attids = []

    # the exact list of attids we send is quite critical. Note that
    # we do ask for the secret attributes, but set set SPECIAL_SECRET_PROCESSING
    # to zero them out
    schema_dn = samdb.get_schema_basedn()
    res = samdb.search(base=schema_dn, scope=ldb.SCOPE_SUBTREE,
                       expression="objectClass=attributeSchema",
                       attrs=["lDAPDisplayName", "systemFlags",
                              "searchFlags"])

    for r in res:
        ldap_display_name = r["lDAPDisplayName"][0]
        if "systemFlags" in r:
            system_flags      = r["systemFlags"][0]
            if (int(system_flags) & (samba.dsdb.DS_FLAG_ATTR_NOT_REPLICATED |
                                     samba.dsdb.DS_FLAG_ATTR_IS_CONSTRUCTED)):
                continue
        if "searchFlags" in r:
            search_flags = r["searchFlags"][0]
            if (int(search_flags) & samba.dsdb.SEARCH_FLAG_RODC_ATTRIBUTE):
                continue
        attid = samdb.get_attid_from_lDAPDisplayName(ldap_display_name)
        attids.append(int(attid))

    # the attids do need to be sorted, or windows doesn't return
    # all the attributes we need
    attids.sort()
    partial_attribute_set.attids         = attids
    partial_attribute_set.num_attids = len(attids)
    return partial_attribute_set


########### main code ###########
if __name__ == "__main__":
    parser = OptionParser("getncchanges [options] server")
    sambaopts = options.SambaOptions(parser)
    parser.add_option_group(sambaopts)
    credopts = options.CredentialsOptionsDouble(parser)
    parser.add_option_group(credopts)

    parser.add_option("", "--dn", dest="dn", help="DN to replicate",)
    parser.add_option("", "--exop", dest="exop", help="extended operation",)
    parser.add_option("", "--pas", dest="use_pas", action='store_true', default=False,
                      help="send partial attribute set (for RODC)")
    parser.add_option("", "--nb-iter", type='int', help="Number of getncchange iterations")
    parser.add_option("", "--dest-dsa", type='str', help="destination DSA GUID")
    parser.add_option("", "--rodc", action='store_true', default=False,
                      help='use RODC replica flags')
    parser.add_option("", "--partial-rw", action='store_true', default=False,
                      help='use RW partial replica flags, not be confused with --pas')
    parser.add_option("", "--replica-flags", type='int',
                      default=drsuapi.DRSUAPI_DRS_INIT_SYNC |
                      drsuapi.DRSUAPI_DRS_PER_SYNC |
                      drsuapi.DRSUAPI_DRS_WRIT_REP |
                      drsuapi.DRSUAPI_DRS_GET_ANC |
                      drsuapi.DRSUAPI_DRS_NEVER_SYNCED,
                      help='replica flags')

    (opts, args) = parser.parse_args()
    if opts.rodc:
        opts.replica_flags = drsuapi.DRSUAPI_DRS_INIT_SYNC |\
                             drsuapi.DRSUAPI_DRS_PER_SYNC |\
                             drsuapi.DRSUAPI_DRS_GET_ANC |\
                             drsuapi.DRSUAPI_DRS_NEVER_SYNCED |\
                             drsuapi.DRSUAPI_DRS_SPECIAL_SECRET_PROCESSING |\
                             drsuapi.DRSUAPI_DRS_GET_ALL_GROUP_MEMBERSHIP

    if opts.partial_rw:
        opts.replica_flags = drsuapi.DRSUAPI_DRS_INIT_SYNC |\
                             drsuapi.DRSUAPI_DRS_PER_SYNC |\
                             drsuapi.DRSUAPI_DRS_GET_ANC |\
                             drsuapi.DRSUAPI_DRS_NEVER_SYNCED

    lp = sambaopts.get_loadparm()
    creds = credopts.get_credentials(lp)

    if len(args) != 1:
        parser.error("You must supply a server")

    if creds.is_anonymous():
        parser.error("You must supply credentials")

    if opts.partial_rw and opts.rodc:
        parser.error("Can't specify --partial-rw and --rodc")

    server = args[0]

    binding_str = "ncacn_ip_tcp:%s[seal,print]" % server

    drs = drsuapi.drsuapi(binding_str, lp, creds)
    drs_handle = do_DsBind(drs)
    print "DRS Handle: %s" % drs_handle

    req8 = drsuapi.DsGetNCChangesRequest8()

    samdb = SamDB(url="ldap://%s" % server,
                  session_info=system_session(),
                  credentials=creds, lp=lp)

    if opts.use_pas:
        local_samdb = SamDB(url=None, session_info=system_session(),
                            credentials=creds, lp=lp)

    if opts.dn is None:
        opts.dn = str(samdb.get_default_basedn())

    if opts.exop is None:
        exop = drsuapi.DRSUAPI_EXOP_NONE
    else:
        exop = int(opts.exop)

    dest_dsa = opts.dest_dsa
    if not dest_dsa:
        print "no dest_dsa specified trying to figure out from ldap"
        msgs = samdb.search(controls=["search_options:1:2"],
                           expression='(objectclass=ntdsdsa)')
        if len(msgs) == 1:
            dest_dsa = str(ndr_unpack(misc.GUID,  msgs[0]["invocationId"][0]))
            print "Found this dsa: %s" % dest_dsa
        else:
            # TODO fixme
            pass
        if not dest_dsa:
            print "Unable to find the dest_dsa automatically please specify it"
            import sys
            sys.exit(1)

    null_guid = misc.GUID()
    req8.destination_dsa_guid               = misc.GUID(dest_dsa)
    req8.source_dsa_invocation_id	    = misc.GUID(samdb.get_invocation_id())
    req8.naming_context			    = drsuapi.DsReplicaObjectIdentifier()
    req8.naming_context.dn                  = opts.dn.decode("utf-8")
    req8.highwatermark                      = drsuapi.DsReplicaHighWaterMark()
    req8.highwatermark.tmp_highest_usn	    = 0
    req8.highwatermark.reserved_usn	    = 0
    req8.highwatermark.highest_usn	    = 0
    req8.uptodateness_vector		    = None
    req8.replica_flags			    = opts.replica_flags
    req8.max_object_count		     = 402
    req8.max_ndr_size			     = 402116
    req8.extended_op			     = exop
    req8.fsmo_info			     = 0
    if opts.use_pas:
        req8.partial_attribute_set	     = drs_get_rodc_partial_attribute_set(local_samdb)
    else:
        req8.partial_attribute_set	     = None
    req8.partial_attribute_set_ex	     = None
    req8.mapping_ctr.num_mappings	     = 0
    req8.mapping_ctr.mappings		     = None

    nb_iter = 0
    while True:
        (level, ctr) = drs.DsGetNCChanges(drs_handle, 8, req8)
        nb_iter += 1
        if ctr.more_data == 0 or opts.nb_iter == nb_iter:
            break
        req8.highwatermark = ctr.new_highwatermark
